import math


def vec_length(v):
    s = 0
    for e in v:
        s += e * e
    if s == 0:
        return 0
    return math.sqrt(s)


def vec_sub(v1, v2):
    result = []
    for i in range(len(v1)):
        result.append(v1[i] - v2[i])
    return result


def distance(v1, v2):
    return vec_length(vec_sub(v1, v2))


def mean(vectors):
    dim = len(vectors[0])
    num_vectors = len(vectors)
    result = []
    for i in range(dim):
        e = 0
        for v in vectors:
            e += v[i]
        e /= num_vectors
        result.append(e)
    return result


def k_mean(samples, k, threshold):
    num_samples = len(samples)
    iter_count = 0

    # 使用前k个样本作为初始聚类中心
    centers = samples[:k]

    while True:
        iter_count += 1
        # 记录每个类别有哪些样本
        class_to_sample = [[] for i in range(k)]
        # 记录每个样本属于哪个类别
        sample_to_class = [-1 for i in range(num_samples)]

        # 计算每个样本属于哪个类别
        for i in range(num_samples):
            sample = samples[i]
            min_dist = 1e+5
            min_dist_idx = -1
            for j in range(k):
                center = centers[j]
                dist = distance(sample, center)
                if dist < min_dist:
                    min_dist = dist
                    min_dist_idx = j
            class_to_sample[min_dist_idx].append(sample)
            sample_to_class[i] = min_dist_idx

        # 重新计算聚类中心
        new_centers = [None] * k
        for i in range(k):
            if len(class_to_sample[i]) == 0:
                new_centers[i] = centers[i]
            else:
                new_centers[i] = mean(class_to_sample[i])

        # 判断是否结束
        is_end = True
        for i in range(k):
            if distance(centers[i], new_centers[i]) > threshold:
                is_end = False
                break
        if is_end:
            return centers, sample_to_class, iter_count
        else:
            centers = new_centers
